from rich.console import Console, Group
from rich.panel import Panel
from rich.columns import Columns
import numpy as np
from dataclasses import dataclass
import inspect
from xarray import DataArray, Dataset
from xarray.core.coordinates import DataArrayCoordinates
from rich.table import Table
from rich.padding import Padding
from rich.tree import Tree
from pathlib import Path

TB = "│"

@dataclass
class ViewerStyle:
    keyword: str = "bright_cyan"
    type: str = "magenta"
    value: str = "default"
    number: str = "green not bold"
    title: str = "blue"
    warning: str = "orange1"
    url: str = "not bold not italic underline bright_blue"

    def __call__(self, style: str, text: str, bold: bool = False) -> str:
        """根据样式名称返回对应的样式
        """
        # 获取自身的属性
        attr = getattr(self, style, None)
        if attr is None:
            raise ValueError(f"Invalid style: {style}")

        if bold:
            return f"[bold {attr}]{text}[/bold {attr}]"
        else:
            return f"[{attr}]{text}[/{attr}]"


class Viewer:
    """
    美化 Numpy 与 Xarray 的打印输出的包装器。

    A warapping class for beautifying the print output of Numpy and Xarray.
    """
    console = None

    def __init__(self, *args, style: ViewerStyle = ViewerStyle(),
                 lineno: bool = True):
        """
        打印任意内容，支持美化 numpy 和 xarray 的打印输出。

        Print any content, supports beautifying the print output of
        numpy and xarray.

        Usage
        -----
        ``` python
        data = np.array([1, 2, 3])

        # 直接使用
        Viewer(data)

        # 或实例化
        v = Viewer()
        v(data)
        ```

        Parameters
        ----------
        *args : any
            需要打印的内容，可以是任意数量的参数。

            The content to be printed, can be any number of parameters.
        style : ViewerStyle, optional
            打印内容的样式，默认为 ViewerStyle()。

            The style to be used for printing, by default ViewerStyle()
        lineno : bool, optional
            是否显示文件与行号，默认为 True。

            Whether to display the file and line number, by default True.
        """
        if self.console is None:
            self.console = Console()

        self.s = style
        self.lineno = lineno
        self.call_frame = inspect.currentframe().f_back

        for arg in args:
            self.console.print(self._print(arg))


    def view(self, *args):
        self.call_frame = inspect.currentframe().f_back
        for arg in args:
            self.console.print(self._print(arg))


    def __call__(self, *args):
        self.call_frame = inspect.currentframe().f_back
        for arg in args:
            self.console.print(self._print(arg))


    def _get_var_name(self, arg) -> str | None:
        """获取变量名"""
        var_name = None
        local_vars = self.call_frame.f_locals
        for name, value in local_vars.items():
            if value is arg:
                var_name = name
                break
        return var_name


    def _get_line_number(self) -> int:
        """获取调用行号"""
        info = inspect.getframeinfo(self.call_frame)
        file_name = Path(info.filename).name
        return f"{file_name}:{info.lineno}"


    @staticmethod
    def _format_bytes(size: int) -> str:
        """格式化字节数
        """
        if size < 1024:
            return f"{size} B"
        elif size < 1024**2:
            return f"{size/1024:.2f} KB"
        elif size < 1024**3:
            return f"{size/1024**2:.2f} MB"
        else:
            return f"{size/1024**3:.2f} GB"


    def _ndarray(self, data: np.ndarray):
        """打印ndarray的简要信息
        """
        info = f"{self.s('number', data.size)} "\
                f"{self.s('type', data.dtype)} values"\
                f" ({self._format_bytes(data.itemsize * data.size)}) "\
                f"@ {self.s('warning', data.shape)}\n" + \
                np.array2string(data, max_line_width=self.console.width-16,
                                edgeitems=3)
        return info


    @staticmethod
    def _datetime64_format(t: np.datetime64) -> str:
        """格式化datetime64
        """
        if t is None:
            return "None"
        elif isinstance(t, np.datetime64):
            return t.astype('datetime64[ms]').item().strftime("%Y-%m-%d %H:%M:%S")
        else:
            return str(t)


    def _xr_coords(self, coord: DataArrayCoordinates, title=True) -> Table:
        """打印xarray的维度信息"""
        table = Table(show_edge=False, show_lines=False,
                      show_header=False, show_footer=False,
                      title_justify='left', title_style=self.s.title)
        table.add_column("index", style=self.s.number + " i")
        table.add_column("name", style=self.s.keyword)
        table.add_column("type", style=self.s.type)
        table.add_column("size", style=self.s.number)
        table.add_column("values",)

        if title:
            table.title = ":compass: Coordinates (index, name, type, size, values)"

        for index, key in enumerate(coord.dims):
            c = coord[key]
            table.add_row(
                str(index), str(key), str(c.dtype), str(c.size),
                np.array2string(c.values, precision=2, suppress_small=True,
                                threshold=4, edgeitems=2, max_line_width=50,
                                formatter={"datetime": self._datetime64_format}
                                ).strip("[]")
            )
        return table

    def _xr_attrs(self, attrs: dict) -> Table:
        """打印属性信息"""
        attrsTable = Table(title= ":memo: Attributes (name, value)",
                            show_edge=False, show_lines=False,
                            show_header=False, show_footer=False,
                            title_style=self.s.title, title_justify="left")

        attrsTable.add_column("name", style=self.s.keyword, justify='right')
        attrsTable.add_column("value", justify='left')

        for key in attrs.keys():
            attrsTable.add_row(key, str(attrs[key]))
        return attrsTable


    def _xr_variable(self, data: DataArray) -> Table:
        """打印xarray变量信息"""
        coords = data.coords
        coords_text = []

        for key in coords.keys():
            coords_text.append(f"{key}: {self.s('number', coords[key].size)}")
        if len(coords_text) > 0:
            coords_text = " (" + ", ".join(coords_text) + ")"
        else:
            coords_text = ""

        tree = Tree(f":package: {self.s('keyword', data.name, True)}{coords_text}",)

        long_name = data.attrs.get("long_name", False)
        std_name = data.attrs.get("standard_name", False)
        description = data.attrs.get("description", False)
        desc = long_name or std_name or description
        unit = data.attrs.get("units", False)
        if desc and unit:
            desc = f"{desc} ({unit})"
        elif desc:
            desc = f"{desc}"
        elif unit:
            desc = f"units: {unit}"
        if desc:
            tree.add(self.s('title', desc))

        size = data.size * data.dtype.itemsize
        compress = data.encoding.get("complevel", 0)
        tree.add(
            f"Mem:{self.s('type', data.dtype)} {self._format_bytes(size)} "
            f"| Encoding: {self.s('type', data.encoding.get('dtype', ''))} "
            f"(complevel: {compress})"
        )
        return tree


    def _print_dataarray(self, data: DataArray, title: str) -> Group:
        source = data.encoding.get("source", "")
        coords = data.coords
        coords_text = []

        for key in coords.keys():
            coords_text.append(f"{key}: {self.s('number', coords[key].size)}")
        if len(coords_text) > 0:
            coords_text = " (" + ", ".join(coords_text) + ")"
        else:
            coords_text = ""

        tree = Tree(f":package: {self.s('keyword', data.name, True)}{coords_text}",)

        size = data.size * data.dtype.itemsize
        compress = data.encoding.get("complevel", 0)
        info = Tree(":information: Information", style=self.s.title)
        info.add(
            f"[default]Mem:{self.s('type', data.dtype)} {self._format_bytes(size)} "
            f"| Encoding: {self.s('type', data.encoding.get('dtype', ''))} "
            f"(complevel: {compress})" +
            (f"\nSource: {self.s('url', source)}" if source else "") + "[/default]"
        )
        tree.add(Padding(info, (0, 0, 1, 0)))
        tree.add(Padding(self._xr_coords(data.coords), (0, 0, 1, 0)))
        tree.add(self._xr_attrs(data.attrs))
        panel = Panel( tree, title=title, style="on black", title_align="left",)
        return panel


    def _print_dataset(self, data: Dataset, title: str, subtitle="") -> Group:
        """打印xarray数据集的简要信息
        """
        source = data.encoding.get("source", "")
        if source != "":
            try:
                file_size = Path(source).stat().st_size
                source = f"{self.s('url', source)} ({self._format_bytes(file_size)})"
            except FileNotFoundError:
                source = f"{self.s('url', source)}"
        else:
            source = "<unknown source or in memory>"
        tree = Tree(f":truck: {source}\n{TB}")

        # variables
        for key in data.data_vars.keys():
            tree.add(Padding(self._xr_variable(data[key]), (0, 0, 1, 0)))

        # dimensions
        tree.add(Padding(self._xr_coords(data.coords), (0, 0, 1, 0)))

        # attributes
        tree.add(self._xr_attrs(data.attrs))

        panel = Panel(
            tree, title=title, style="on black",
            subtitle=subtitle, title_align="left",
        )
        return panel


    def _print(self, v: any):
        """打印信息"""
        name = self._get_var_name(v)
        var_title = ""
        if self.lineno:
            var_title += f"{self.s('number', self._get_line_number())} "
        var_title += f"{self.s('type', f'<{type(v).__name__}>')}"

        if name is not None:
            var_title += f" {self.s('keyword', self._get_var_name(v), True)}:"

        if isinstance(v, np.ndarray):
            var_info = self._ndarray(v)
            return Columns([var_title, var_info])
        elif isinstance(v, DataArray):
            return self._print_dataarray(v, var_title.strip(":"))
        elif isinstance(v, Dataset):
            return self._print_dataset(v, var_title.strip(":"))
        else:
            var_info = str(v)
            return Columns([var_title, var_info])


    def _insert_str(origin: str, insertList: list[str],
                    insertIndexList: list[int]) -> str:
        """插入字符串
        """
        for i in range(len(insertList)):
            strLen = len(insertList[i])
            if insertIndexList[i] == 0:
                origin = insertList[i] + origin[strLen:]
            elif insertIndexList[i] == len(origin) - 1:
                origin = origin[:-strLen] + insertList[i]
            else:
                halfLen = strLen // 2
                if strLen % 2 == 0:
                    origin = origin[:insertIndexList[i]-halfLen] + \
                        insertList[i] + origin[insertIndexList[i]+halfLen:]
                else:
                    origin = origin[:insertIndexList[i]-halfLen] + \
                        insertList[i] + origin[insertIndexList[i]+halfLen+1:]
        return origin


    @staticmethod
    def pcolor(data: DataArray | np.ndarray, x: DataArray | np.ndarray = None,
            y: DataArray | np.ndarray = None, x_name: str = "longitude",
            y_name: str = "latitude", cmap=None, coorTickNum=5,
            colorbarTickNum=5, zeroCenter=None):
        """在终端中显示2D数据的伪彩色图。

        Display a pseudo-color map of 2D data in the terminal.

        Parameters
        ----------
        data : DataArray | np.ndarray
            数据，需要为2维。

            The data to be displayed, must be 2D.
        x, y : DataArray | np.ndarray, optional
            指定数据x 轴和 y 轴的坐标。

            Specify the coordinates of the x-axis and y-axis of the data.
        x_name, y_name: str, optional
            指定 DataArray 作为 x 轴和 y 轴的坐标的名称。

            Specify the name of the DataArray as the coordinates of
            the x-axis and y-axis.
        cmap : _type_, optional
            颜色的 Colormap，默认为 None，如果为 None，则根据数据自动选择
            `turbo` 或 `RdBu_r`。

            The Colormap of the color, default is None, if None,
            it will be automatically selected according to the data
            `turbo` or `RdBu_r`.
        coorTickNum : int, optional
            坐标刻度的数量，默认为 5。

            The number of coordinate ticks, default is 5.
        colorbarTickNum : int, optional
            颜色条刻度的数量，默认为 5。

            The number of color bar ticks, default is 5.
        zeroCenter : _type_, optional
            颜色值是否以 0 为中心，默认为 None，如果为 None，则根据数据自动选择。

            Whether the color value is centered at 0, default is None,
            if None, it will be automatically selected according to the data.

        """
        if isinstance(data, (list, tuple)):
            try:
                data = np.array(data)
            except Exception as e:
                raise TypeError("failed to convert data to ndarray") from e

        # 如果是DataArray，检查是否有坐标
        if isinstance(data, DataArray):
            # 如果直接指定
            if x is not None and y is not None:
                x = x.values if isinstance(x, DataArray) else x
                y = y.values if isinstance(y, DataArray) else y

            # 没有指定，检查是否有坐标
            elif hasattr(data, x_name) and hasattr(data, y_name):
                x = data[x_name].values
                y = data[y_name].values

            else:
                x = data.coords[data.dims[-1]].values
                y = data.coords[data.dims[-2]].values

            name = data.name
            data = data.values

        elif isinstance(data, np.ndarray):
            if x is None:
                x = np.arange(data.shape[1])
            if y is None:
                y = np.arange(data.shape[0])
            name = "ndarray"
        else:
            raise TypeError("input data must be DataArray, ndarray, or 2D list/tuple"
                            " that can be converted to ndarray")


        # 检查数据是否为2D
        if len(data.shape) != 2:
            raise ValueError(f"只能显示2维数据，当前数据 shape 为{len(data.shape)}")

        # 检查坐标是否匹配
        if len(x) != data.shape[1] or len(y) != data.shape[0]:
            raise ValueError("the length of x and y must match the shape of data")

        # 终端画图
        from matplotlib.cm import get_cmap
        if Viewer.console is None:
            Viewer.console = Console()

        # 宽高缩放
        width = Viewer.console.width - 10
        if width > data.shape[1]:
            width = data.shape[1]
            height = data.shape[0]
            paintData = data
        else:
            height = int(data.shape[0]/data.shape[1] * width)
            xIndex = np.linspace(0, data.shape[1], width, dtype=int,
                                 endpoint=False)
            yIndex = np.linspace(0, data.shape[0], height, dtype=int,
                                 endpoint=False)
            paintData = data[yIndex, :][:, xIndex]

        # 数值缩放
        text = ""
        maxVal = np.nanmax( paintData )
        minVal = np.nanmin( paintData )
        realMax = maxVal.copy()
        realMin = minVal.copy()
        amp = int(np.log10(np.nanmax(np.abs([maxVal, minVal]))))

        if zeroCenter is None:
            # 自动识别是否需要以0为中心
            if minVal < 0 and maxVal > 0 and \
                np.abs(maxVal+minVal) / ((maxVal-minVal) / 2) < 0.3:
                zeroCenter = True

        if zeroCenter:
            theMax = np.max([np.abs(maxVal), np.abs(minVal)])
            maxVal = theMax
            minVal = -theMax

        if cmap is None:
            if zeroCenter:
                cmap = "RdBu_r"
            else:
                cmap = "turbo"

        cmap = get_cmap(cmap)

        paintData = (paintData - minVal) / (maxVal - minVal)

        # ytick
        ytickLoc = np.linspace(0, height, coorTickNum)
        ytickValue = y[::(data.shape[0]//(coorTickNum-1))]
        ytickValue = [f"{i}" for i in ytickValue]
        ytickLoc = [int(i//2) for i in ytickLoc]

        # 画图
        for i in range(height//2):
            for j in range(width):
                c = cmap(paintData[i*2, j])[:3]
                c = f"rgb({int(c[0]*255)},{int(c[1]*255)},{int(c[2]*255)})"
                c1 = cmap(paintData[i*2+1, j])[:3]
                c1 = f"rgb({int(c1[0]*255)},{int(c1[1]*255)},{int(c1[2]*255)})"
                text += f"[{c} on {c1}]▀[/{c} on {c1}]"
            if i in ytickLoc:
                text += "─" + ytickValue[ytickLoc.index(i)]
            elif i == height//2-1:
                text += "─" + ytickValue[-1]
            if i != height//2-1:
                text += "\n"

        # xtick
        tickLoc = np.linspace(0, width, coorTickNum)
        tickValue = x[::(data.shape[1]//(coorTickNum-1))]
        tickValue = [f"{i}" for i in tickValue]
        tickLoc = [int(i) for i in tickLoc]
        tickText = " " * width
        tickText = Viewer._insert_str(tickText, "│"*coorTickNum, tickLoc)
        tickValueText = " " * width
        tickValueText = Viewer._insert_str(tickValueText, tickValue, tickLoc)
        text += "\n" + tickText + "\n" + tickValueText + "\n"

        # colorbar
        colorbarValue = np.linspace(minVal, maxVal, width)
        colorbarText = ""
        for i in range(width):
            c = cmap((colorbarValue[i] - minVal) / (maxVal - minVal))[:3]
            c = f"rgb({int(c[0]*255)},{int(c[1]*255)},{int(c[2]*255)})"
            colorbarText += f"[{c}]█[/{c}]"

        # colorbar tick
        tickLoc = np.linspace(0, width, colorbarTickNum)
        tickLoc = [int(i) for i in tickLoc]

        tickValue = np.linspace(minVal*(10**-amp), maxVal*(10**-amp), colorbarTickNum)
        tickValue = [f"{i:.2f}" for i in tickValue]
        tickValueText = " " * width
        tickValueText = Viewer._insert_str(tickValueText, tickValue, tickLoc)
        tickText = "─" * width
        tickText = Viewer._insert_str(tickText, "┌"+("┬"*(colorbarTickNum-2))+"┐", tickLoc)

        # 说明
        text += "\n" + colorbarText + "\n" + tickText + "\n" + tickValueText + "\n"
        text += f"values = [blue]color × 10^{amp}[/blue]  max:[blue]{realMax:.3g}[/blue]  min:[blue]{realMin:.3g}[/blue]"

        # 显示
        zoom = width / data.shape[1]
        panel = Panel.fit(text, title=f"View of {name} "
                        f"({data.shape[1]}*{data.shape[0]} zoom: {zoom:.2f})",
                        style="on black", subtitle=f"({width}*{height}px)")
        Viewer.console.print(panel)



